import pandas as pd
import matplotlib.pyplot as plt
import okx.MarketData as MarketData

flag = "0"
marketDataAPI =  MarketData.MarketAPI(flag=flag)
result = marketDataAPI.get_candlesticks(instId="BTC-USDT",bar = '1H')
raw_data = result["data"]

# 创建DataFrame
df = pd.DataFrame(raw_data, columns=[
    'timestamp', 'open', 'high', 'low', 'close', 'vol', 'volCcy', 'volCcyQuote', 'confirm'
])
df['Date'] = pd.to_datetime(df['timestamp'].astype(int), unit='ms')
df.set_index('Date', inplace=True)
# 转换数值型
for col in ['open', 'high', 'low', 'close', 'vol']:
    df[col] = df[col].astype(float)
df = df.sort_index()  # 按时间升序

# -----------------
# KDJ计算部分
# -----------------
n = 9  # 常用周期，样本太短建议设成比样本长度小的值(如3)
low_list = df['low'].rolling(window=n, min_periods=1).min()
high_list = df['high'].rolling(window=n, min_periods=1).max()
rsv = (df['close'] - low_list) / (high_list - low_list) * 100

df['K'] = rsv.ewm(com=2, adjust=False).mean()  # K线
df['D'] = df['K'].ewm(com=2, adjust=False).mean()  # D线
df['J'] = 3 * df['K'] - 2 * df['D']  # J线

# -----------------
# 绘图保存
# -----------------
plt.figure(figsize=(10,5))
plt.plot(df.index, df['K'], label='K', color='Blue')
plt.plot(df.index, df['D'], label='D', color='Green')
plt.plot(df.index, df['J'], label='J', color='Red')
plt.legend()
plt.title('KDJ指标')
plt.xlabel('时间')
plt.ylabel('数值')
plt.tight_layout()
plt.savefig('crypto_kdj.png')
plt.close()
print('KDJ图已保存到本地：crypto_kdj.png')
